"""定义系统的odes"""
import numpy as np
from math import sin, pi, floor

# 为灵活使用以及避免每次都计算某些参数，定义描述系统的类，参数使用实例的属性来记录
class TIA():
    """
    描述差分形式跨阻放大器之参数和动态过程的类
    两个运放均假设为一阶系统
    """
    def __init__(self, A0=[1e6, 1e6], bw=1.6e7, Rf=[1e7, 1e7], \
        Cf=[3.3e-12, 3.3e-12], Cin=[1e-14, 1e-14], en=None):
        """
        初始化差分跨阻放大器的参数
        参数：
        A0：两个运放的开环增益
        bw：两个运放的带宽
        Rf：两个运放的反馈电阻
        Cf：两个运放的反馈电容
        Cin：两个运放的输入电容
        en：两个运放的等效输入电压噪声
        """
        # 参数
        self.A0 = np.array(A0, dtype=np.float64)
        self.tao = self.A0/(2*pi*bw)
        self.Rf = np.array(Rf, dtype=np.float64)
        self.Cf = np.array(Cf, dtype=np.float64)
        self.Cin = np.array(Cin, dtype=np.float64)
        # 噪声
        self.en = en
        # 系统矩阵中的元素
        temp = self.Rf*(self.Cin + self.Cf)*self.tao
        b_1 = (self.Rf*self.Cf*(1+self.A0) + self.tao + self.Rf*self.Cin)/temp
        b_0 = (self.A0 + 1)/temp
        x_1Up = self.A0*self.Rf*self.Cf/temp
        x_1Upe = self.A0*self.Rf*(self.Cin + self.Cf)/temp
        x_0Up = self.A0/temp - self.A0*self.Rf*self.Cf/temp*b_1
        x_0Upe = self.A0/temp - self.A0*self.Rf*(self.Cf + self.Cin)/temp*b_1
        x_0I = -self.A0*self.Rf/temp
        # 系统矩阵
        self.A = np.empty((2,2,2), dtype=np.float64)
        self.BI = np.empty((2,2), dtype=np.float64)
        self.BUp = np.empty((2,2), dtype=np.float64)
        self.BUpe = np.empty((2,2), dtype=np.float64)
        self.A[0,:,:] = np.array([[0,1], [-b_0[0],-b_1[0]]], dtype=np.float64)
        self.A[1,:,:] = np.array([[0,1], [-b_0[1],-b_1[1]]], dtype=np.float64)
        self.BI[0,:] = np.array([0, x_0I[0]], dtype=np.float64)
        self.BI[1,:] = np.array([0, x_0I[1]], dtype=np.float64)
        self.BUp[0,:] = np.array([x_1Up[0], x_0Up[0]], dtype=np.float64)
        self.BUp[1,:] = np.array([x_1Up[1], x_0Up[1]], dtype=np.float64)
        self.BUpe[0,:] = np.array([x_1Upe[0], x_0Upe[0]], dtype=np.float64)
        self.BUpe[1,:] = np.array([x_1Upe[1], x_0Upe[1]], dtype=np.float64)

    def TIA_odes(self, Is, OP_states, noise_index=None):
        """
        描述差分跨阻放大器的动态过程的微分方程
        参数：
        Is：输入电流
        OP_states：运放状态（输出及一阶）
        noise_index：噪声矩阵的索引
        """
        TIA_states = np.empty((4,), dtype=np.float64)
        # 等效输入电压噪声、等效输入电流噪声、反馈阻抗热噪声
        # Upe = en[noise_index, :]
        # 求运放状态的导数
        TIA_states[0:2] = self.A[0,:,:]@OP_states[0:2]
        TIA_states[0:2] += self.BI[0,:]*(-Is)
        # TIA_states[0:2] += self.BUpe[0,:]*Upe[0]
        # TIA_states[0:2] += self.BUp[0,:]*Up[0]
        TIA_states[2:] = self.A[1,:,:]@OP_states[2:]
        TIA_states[2:] += self.BI[1,:]*(Is)
        # TIA_states[2:] += self.BUpe[1,:]*Upe[1]
        # TIA_states[2:] += self.BUp[1,:]*Up[1]
        # 求负输入端的电压
        # Un = -(OP_states[[0,2]]+self.tao*OP_states[[1,3]])/self.A0没啥性能优势
        Un = [
            # Up[0]+Upe[0]-(OP_states[0]+self.tao[0]*OP_states[1])/self.A0[0],
            # Up[1]+Upe[1]-(OP_states[2]+self.tao[1]*OP_states[3])/self.A0[1]
            -(OP_states[0]+self.tao[0]*OP_states[1])/self.A0[0],
            -(OP_states[2]+self.tao[1]*OP_states[3])/self.A0[1]
        ]
        return TIA_states, Un


class Bridge():
    """描述3绕组（2个一次侧和1个二次侧）差分变压器桥之参数和动态过程的类"""
    def __init__(self, switching, C_pi, C_p, C_sense=[1.15e-12,1.15e-12], 
            L_p1=5e-3, L_s=5e-3, L_p2=5e-3, r_l=[31.4, 31.4, 31.4], 
            k12=0.99, k13=0.99, k23=0.99, sig_freq=0.1, pump_freq=1e5, 
            tao=[70, 350], gain=25, U_m=0.6, th_noise=None):
        """初始化差分变压器桥的参数"""
        # 自感
        self.L_p1 = L_p1
        self.L_p2 = L_p2
        self.L_s = L_s
        # 敏感电容、补偿电容阵列和主调谐电容
        self.C_sense = np.array(C_sense, dtype=np.float64)
        self.C_pi = C_pi
        self.C_p = C_p
        # 互感
        self.M_12 = k12*(L_p1*L_p2)**0.5
        self.M_21 = self.M_12# 假设变压器对称
        self.M_13 = k13*(L_p1*L_s)**0.5
        self.M_31 = self.M_13
        self.M_23 = k23*(L_p2*L_s)**0.5
        self.M_32 = self.M_23
        # 绕组损失(电阻)
        self.r_l = np.array(r_l, dtype=np.float64)
        # 补偿电容阵列开关（电阻）
        self.switching = switching
        self.nswitch = int(0.5*switching.size)
        self.Yswh = -1/self.switching
        self.YswhCpi = self.Yswh/self.C_pi
        # 信号
        self.U_m = np.float64(U_m)
        self.omega_C = np.float64(2*pi*sig_freq)# C1和C2变化的角频率
        self.omega = np.float64(2*pi*pump_freq)# 激励频率
        # 求解电感电流的导数所需的矩阵
        L_M = np.array(
            [[L_p1, -self.M_12, self.M_13],
            [-self.M_21, L_p2, -self.M_23],
            [self.M_31, -self.M_32, L_s]], dtype=np.float64)
        self.e_factors = np.linalg.inv(L_M)
        # TIA的odes
        self.tia = TIA()
        # 低通滤波器的时间常数
        tao[0] *= 2*pi
        tao[1] *= 2*pi
        self.lp_factor = np.array([tao[0]*tao[1], tao[0]+tao[1]],\
            dtype=np.float64)
        # TIA之后所有的增益(解调将导致增益*0.5)
        self.gain = np.float64(gain)*2
        # 电感损失电阻的热噪声
        self.th_noise = th_noise
        # 暂存雅可比矩阵
        jac_matrix = self.jacobian_full(0, \
            np.zeros((11+2*self.nswitch,),dtype=np.float64))
        self.jac_matrix = jac_matrix

    def bridge_odes(self, t, states):
        """描述桥路动态过程的微分方程
        若使用odeint()，注意tfirst=True
        """
        dstates_dt = np.empty(states.shape, dtype=np.float64)
        # 激励信号
        temp = sin(self.omega*t)# np设施返回的是double
        U_m = self.U_m*temp
        # 中间变量
        demodulate_U_o = states[-6] - states[-4]
        demodulate_U_o *= temp
        demodulate_U_o *= self.gain
        temp = 6e-14*sin(self.omega_C*t)
        C_sense = self.C_sense + np.array([temp, -temp], dtype=np.float64)
        C_eq = C_sense + self.C_p
        U_p = C_sense/C_eq*U_m
        # U_ceq = states[3:5]/C_eq# 等效电容Ceq//C_p上的电压
        U_p -= states[3:5]/C_eq
        # TIA 状态导数
        # noise_index = floor((t - floor(t))/5e-6)
        dstates_dt[-6:-2], Un = self.tia.TIA_odes(states[2], states[-6:-2])
        # 桥路电感状态（电流）的导数
        temp = np.array( [U_p[0], U_p[1], Un[0]-Un[1]])
        temp -= states[0:3]*self.r_l
        dstates_dt[0:3] = self.e_factors@temp
        # 桥路电容状态（电荷）的导数
        U_prii = states[5:-6]/self.C_pi# 电容阵列中各路电容的电压
        temp = np.array([U_p[0], U_p[0], U_p[0], U_p[1], U_p[1], U_p[1]])
        dstates_dt[5:-6] = (temp - U_prii)/self.switching
        half = 5 + self.nswitch
        dstates_dt[3:5] = [np.sum(dstates_dt[5:half]), \
            np.sum(dstates_dt[half:-6])]
        dstates_dt[3:5] += states[0:2]
        # 低通滤波器
        dstates_dt[-2:] = np.array([states[-1],\
            self.lp_factor[0]*demodulate_U_o - self.lp_factor@states[-2:]],\
            dtype=np.float64)
        return dstates_dt

    def jacobian_full(self, t, states):
        """
        定义系统状态方程的雅克比矩阵（初始化）。
        由于该矩阵中有许多不会变化的元素，所以只在初始化时运行完整计算过程
        """
        df_dy = np.zeros((states.size, states.size), dtype=np.float64)
        # 中间变量
        temp = self.omega_C*t
        temp = sin(temp)
        temp *= 6e-14
        C_sense = self.C_sense + np.array( [temp, -temp], dtype=np.float64)
        C_eq = C_sense + self.C_p
        temp = np.empty((4,), dtype=np.float64)
        temp1 = 1/self.tia.A0
        temp[[0,2]] = -temp1[0], temp1[1]
        temp1 *= self.tia.tao
        temp[[1,3]] = -temp1[0], temp1[1]
        half = 5 + self.nswitch
        # temp = np.array([-temp1[0], -self.tia.tao[0]/self.tia.A0[0],\
        #     temp1[1], self.tia.tao[1]/self.tia.A0[1]],dtype=np.float64)
        # df_dIp1
        df_dy[0,0:3] = -self.e_factors[0,:]*self.r_l
        df_dy[0,3:5] = -self.e_factors[0,0:2]/C_eq
        df_dy[0,-6:-2] = self.e_factors[0,2]*temp
        # df_dIp2
        df_dy[1,0:3] = -self.e_factors[1,:]*self.r_l
        df_dy[1,3:5] = -self.e_factors[1,0:2]/C_eq
        df_dy[1,-6:-2] = self.e_factors[1,2]*temp
        # df_dIs
        df_dy[2,0:3] = -self.e_factors[2,:]*self.r_l
        df_dy[2,3:5] = -self.e_factors[2,0:2]/C_eq
        df_dy[2,-6:-2] = self.e_factors[2,2]*temp
        # df_dQceq1 & df_dQceq2
        df_dy[3,0] = 1
        temp = -1/self.switching
        # df_dy[3,3] = -1*np.sum(temp1[0:self.nswitch])/C_eq[0]
        df_dy[3,3] = np.sum(temp[0:self.nswitch])
        df_dy[3,3] /= C_eq[0]
        df_dy[4,1] = 1
        df_dy[4,4] = np.sum(temp[self.nswitch:])
        df_dy[4,4] /= C_eq[1]
        temp1 = temp/self.C_pi
        df_dy[3,5:half] = temp1[0:self.nswitch]     
        df_dy[4,half:(half+self.nswitch)] = temp1[self.nswitch:]
        # df_dQpri
        # temp = -1/self.switching
        df_dy[5:half,3] = temp[0:self.nswitch]/C_eq[0]
        df_dy[half:(half+self.nswitch),4] = temp[self.nswitch:]/C_eq[1]
        # temp /= self.C_pi
        temp_index = np.arange(5, half+self.nswitch,dtype=int)
        df_dy[temp_index,temp_index] = temp1
        # df_dtia状态
        df_dy[-5, 2] = -self.tia.BI[0,1]
        df_dy[-6:-4,-6:-4] = self.tia.A[0,:,:]
        df_dy[-3, 2] = self.tia.BI[1,1]
        df_dy[-4:-2,-4:-2] = self.tia.A[1,:,:]
        # df_d低通滤波器状态
        df_dy[-2,-1] = 1
        df_dy[-1,-2:] = -self.lp_factor
        temp = self.omega*t
        temp = sin(temp)
        temp *= self.lp_factor[0]
        temp *= self.gain# temp = self.lp_factor[0]*sin(self.omega*t)
        df_dy[-1,-6] = temp
        df_dy[-1,-4] = -temp
        return df_dy

    def system_switching(self, switching):
        """投切开关时系统相关参数的变化"""
        # 开关状态变化
        self.switching = switching
        self.Yswh = -1/self.switching# 谐振控制时，switching时变
        self.YswhCpi = self.Yswh/self.C_pi
        half = 5 + self.nswitch
        # df_dQceq1 & df_dQceq2
        self.jac_matrix[3,5:half] = self.YswhCpi[0:self.nswitch]
        self.jac_matrix[4,half:-6] = self.YswhCpi[self.nswitch:]
        # df_dQpri
        temp_index = np.arange(5, half+self.nswitch,dtype=int)
        self.jac_matrix[temp_index,temp_index] = self.YswhCpi

    def jacobian(self, t, states):
        """雅克比矩阵计算（正常运行时）"""
        # 中间变量
        temp = self.omega_C*t
        temp = sin(temp)
        temp *= 6e-14
        C_sense = self.C_sense + np.array( [temp, -temp], dtype=np.float64)
        C_eq = C_sense + self.C_p
        half = 5 + self.nswitch
        # df_dIp1 df_dIp2 df_dIs
        self.jac_matrix[:3,3:5] = -self.e_factors[:,0:2]/C_eq
        # df_dQceq1 & df_dQceq2
        self.jac_matrix[[3,4],[3,4]] = \
            np.sum(self.Yswh[0:self.nswitch]), np.sum(self.Yswh[self.nswitch:])
        self.jac_matrix[[3,4],[3,4]] /= C_eq
        # df_dQpri
        self.jac_matrix[5:half,3] = self.Yswh[0:self.nswitch]/C_eq[0]
        self.jac_matrix[half:-6,4] = self.Yswh[self.nswitch:]/C_eq[1]
        # 低通滤波器状态
        temp = sin(self.omega*t)
        temp *= self.lp_factor[0]
        temp *= self.gain
        self.jac_matrix[-1,-6] = temp
        self.jac_matrix[-1,-4] = -temp
        return self.jac_matrix
